import numpy as np
import random
import os
import gym
import gym.spaces as spaces
from typing import Optional
from itertools import cycle

import cv2
import pygame


def get_hit_mask(image):
    """returns a hitmask using an image's alpha."""
    mask = []
    for x in range(image.get_width()):
        mask.append([])
        for y in range(image.get_height()):
            mask[x].append(bool(image.get_at((x, y))[3]))
    return mask


class GameWindow(object):
    """游戏窗口"""

    def __init__(self):
        self.SCREEN_W, self.SCREEN_H = 276, 512
        self.SCREEN = pygame.display.set_mode((self.SCREEN_W, self.SCREEN_H))
        pygame.display.set_caption("Flappy bird")
        self.BACKGROUND_COLOR = (77, 192, 202)

    def fill_background(self):
        self.SCREEN.fill(self.BACKGROUND_COLOR)


class Base:
    SPRITES = os.path.join(
        os.path.abspath(os.path.dirname(__file__)),
        os.path.join('assets', 'sprites')
    )
    VEL_X = -4  # 管道 X 轴移动速度

    def __init__(self, img_paths):
        if type(img_paths) != list:
            img_paths = [img_paths]
        self.IMAGES = [pygame.image.load(os.path.join(Base.SPRITES, path)).convert_alpha() for path in img_paths]
        self.WIDTH, self.HEIGHT = self.IMAGES[0].get_width(), self.IMAGES[0].get_height()
        self.MASK = [get_hit_mask(i) for i in self.IMAGES]
        self.x, self.y = 0, 0

    def draw(self, screen):
        screen.blit(self.IMAGES[0], (self.x, self.y))


class Bird(Base):
    MAX_VEL_Y = 10  # 沿 Y 的最大（下降）速度
    MIN_VEL_Y = -8  # 沿 Y 的最小速度，最大上升速度
    FLAP_VEL = -7  # 鸟飞时的向上的初速度
    ACC_Y = 1  # 鸟的向下加速度
    INDEX_GEN = cycle([0, 1, 2, 1])  # 改变小鸟形状
    IMG_PATHS = [
        'redbird-up-flap.png',
        'redbird-mid-flap.png',
        'redbird-down-flap.png'
    ]

    def __init__(self, x, y):
        super().__init__(Bird.IMG_PATHS)
        self.x, self.y = x, y
        self.MID_X = self.x + self.WIDTH / 2  # 鸟的中心
        self.vel_y = 0  # 鸟沿 Y 轴的速度
        self.img_index = self.loop_iter = 0

    def move(self, action, base_y):
        if action == FlappyBirdEnv.FLAP:
            if self.y > -2 * self.HEIGHT:  # 鸟停留在可飞范围，可以上飞
                self.vel_y = self.FLAP_VEL  # 给（向上的）一个初速度

        if self.vel_y < self.MAX_VEL_Y and action is FlappyBirdEnv.DO_NOTHING:
            self.vel_y += self.ACC_Y  # 此刻没飞，给个（向下的）加速度
        self.y += min(self.vel_y, base_y - self.y - self.HEIGHT)  # 鸟到地面的距离，确保不钻地
        if self.y < 0:  # 确保不升天
            self.y = 0

        # 鸟的形状改变
        if (self.loop_iter + 1) % 3 == 0:
            self.img_index = next(Bird.INDEX_GEN)
        self.loop_iter = (self.loop_iter + 1) % 30  # 确保 loopIter 不会无限增加

    def draw(self, screen):
        screen.blit(self.IMAGES[self.img_index], (self.x, self.y))


class Ground(Base):
    IMG_PATHS = 'base.png'

    def __init__(self, screen_w, screen_h):
        super().__init__(Ground.IMG_PATHS)
        self.x, self.y = 0, screen_h * 0.79
        self.SHIFT = self.WIDTH - screen_w  # 地面平移
        # 为了让地面平移时看起来更顺滑，shift 必须要头尾接上
        self.SHIFT = self.SHIFT - self.SHIFT % (self.WIDTH / 14)

    def move(self):
        self.x = (self.x + self.VEL_X) % (-self.SHIFT)  # 地面后移


class Pipe(Base):
    IMG_PATHS = 'pipe-green.png'
    GAP_SIZE = 100  # 管道上下间隙
    GAP_YS = [20, 30, 40, 50, 60, 70, 80, 90]
    GAP_X = 200  # 两管道间距

    def __init__(self, x):
        super().__init__(Pipe.IMG_PATHS)
        self.x = x

    def move(self):
        self.x += self.VEL_X


class UpperPipe(Pipe):
    def __init__(self, x, gap_y):
        super(UpperPipe, self).__init__(x)
        self.y = gap_y - self.HEIGHT
        self.IMAGES[0] = pygame.transform.rotate(self.IMAGES[0], 180)
        self.MASK[0] = get_hit_mask(self.IMAGES[0])


class LowerPipe(Pipe):
    def __init__(self, x, gap_y):
        super(LowerPipe, self).__init__(x)
        self.y = gap_y + Pipe.GAP_SIZE


class Score(Base):
    IMG_PATHS = ['{}.png'.format(i) for i in range(10)]
    MARGIN = 2

    def __init__(self, screen_w, screen_h):
        super(Score, self).__init__(Score.IMG_PATHS)
        self.score = 0
        self.x_offset = screen_w - Score.MARGIN
        self.y_offset = screen_h - Score.MARGIN

    def draw(self, screen):
        """displays score in corner of screen"""
        x = self.x_offset

        for digit in reversed([int(i) for i in list(str(self.score))]):
            y = self.y_offset - self.IMAGES[digit].get_height()
            x -= self.IMAGES[digit].get_width()
            screen.blit(self.IMAGES[digit], (x, y))


class FlappyBirdEnv(gym.core.Env, GameWindow):
    metadata = {'render.modes': ['human']}
    reward_range = (-1000, 1)
    spec = None
    DO_NOTHING, FLAP = 0, 1  # actions

    def __init__(self):
        GameWindow.__init__(self)
        pygame.init()  # 初始化所有pygame模块

        self.background = Base('sky.png')
        if self.SCREEN_W != self.background.WIDTH:
            self.SCREEN_W = self.background.WIDTH
            self.SCREEN = pygame.display.set_mode((self.SCREEN_W, self.SCREEN_H))  # 调整窗口

        self.ground = Ground(self.SCREEN_W, self.SCREEN_H)
        self.bird = Bird(int(self.SCREEN_W * 0.2), int((self.SCREEN_H - self.ground.HEIGHT) / 2))
        u_pipe, l_pipe = self.get_random_pipe()
        self.upper_pipes, self.lower_pipes = [u_pipe], [l_pipe]  # 游戏管道

        self.background.y = self.ground.y - self.ground.HEIGHT + 5
        self.score = Score(self.SCREEN_W, self.SCREEN_H)

        # 环境参数
        self.action_space = spaces.Discrete(2)
        self.OBS_W, self.OBS_H = self.SCREEN_W // 2, self.SCREEN_H // 2
        # self.OBS_W, self.OBS_H = 80, 80
        self.observation_space = spaces.Box(
            low=0, high=255, shape=(4, self.OBS_W, self.OBS_H), dtype=np.int
        )
        self.old_obs: Optional[np.ndarray] = None

    def step(self, action: int):
        pygame.event.pump()  # 内部处理pygame事件处理程序

        self.bird.move(action, self.ground.y)
        self.ground.move()
        # 管道左移
        for u_pipe, l_pipe in zip(self.upper_pipes, self.lower_pipes):
            u_pipe.move(), l_pipe.move()

        # 当第一个管道即将触摸屏幕左侧时，添加新管道
        if 0 < self.upper_pipes[0].x - (self.SCREEN_W - Pipe.GAP_X) <= abs(Pipe.VEL_X):
            new_pipe = self.get_random_pipe()
            self.upper_pipes.append(new_pipe[0]), self.lower_pipes.append(new_pipe[1])

        # 如果第一个管道不在屏幕上，则将其删除
        if self.upper_pipes[0].x < -self.upper_pipes[0].WIDTH:
            self.upper_pipes.pop(0), self.lower_pipes.pop(0)

        done = self.check_collision()
        reward = -1000 if done else self.cal_reward()
        # observation = self.process_frame(self.update_frame())  # obs预处理
        observation = self.update_frame()  # obs预处理

        # 如果你想看到训练效果，取消下面的注释
        # self.render()

        return observation, reward, done, {}

    def reset(self):
        self.__init__()
        return self.step(FlappyBirdEnv.DO_NOTHING)[0]

    def cal_reward(self):
        reward = 1

        for pipe in self.upper_pipes:
            pipe_mid_pos = pipe.x + pipe.WIDTH / 2
            if pipe_mid_pos <= self.bird.MID_X < pipe_mid_pos + abs(Pipe.VEL_X):
                self.score.score += 1
                reward = 1
                break

        return reward

    def update_frame(self):
        self.fill_background()
        self.background.draw(self.SCREEN)

        for u_pipe, l_pipe in zip(self.upper_pipes, self.lower_pipes):
            u_pipe.draw(self.SCREEN), l_pipe.draw(self.SCREEN)

        self.ground.draw(self.SCREEN)
        self.bird.draw(self.SCREEN)
        self.score.draw(self.SCREEN)

        return pygame.surfarray.array3d(pygame.display.get_surface())

    def process_frame(self, frame):
        """obs预处理"""
        frame = cv2.cvtColor(cv2.resize(frame, (self.OBS_H, self.OBS_W)), cv2.COLOR_BGR2GRAY)  # 灰度化
        _, frame = cv2.threshold(frame, 150, 255, cv2.THRESH_BINARY)  # 阈值化

        if type(self.old_obs) is not np.ndarray:
            frame = np.reshape(frame, (self.OBS_W, self.OBS_H))
            frame = np.stack((frame, frame, frame, frame), axis=0)  # old_obs 为 None，使用4个相同的 obs 叠加
            self.old_obs = frame
        else:
            frame = np.reshape(frame, (1, self.OBS_W, self.OBS_H))
            self.old_obs = np.append(frame, self.old_obs[:3, :, :], axis=0)  # 使用当前的 obs 和 old_obs 叠加

        return self.old_obs

    def get_random_pipe(self):
        """returns a randomly generated pipe"""
        # y of gap between upper and lower pipe
        gap_y = random.choice(Pipe.GAP_YS) + int(self.ground.HEIGHT * 0.2)
        pipe_x = self.SCREEN_W

        return UpperPipe(pipe_x, gap_y), LowerPipe(pipe_x, gap_y)

    def check_collision(self):
        """returns True if bird collides with base or pipes."""

        # if player crashes into ground
        if self.bird.y + self.bird.HEIGHT >= self.ground.y - 1:
            return True
        else:
            bird_rect = pygame.Rect(self.bird.x, self.bird.y,
                                    self.bird.WIDTH, self.bird.HEIGHT)
            for u_pipe, l_pipe in zip(self.upper_pipes, self.lower_pipes):
                # upper and lower pipe rects
                u_pipe_rect = pygame.Rect(u_pipe.x, u_pipe.y, u_pipe.WIDTH, u_pipe.HEIGHT)
                l_pipe_rect = pygame.Rect(l_pipe.x, l_pipe.y, l_pipe.WIDTH, l_pipe.HEIGHT)

                # player and upper/lower pipe hitmasks
                p_hit_mask = self.bird.MASK[self.bird.img_index]
                u_hit_mask = self.upper_pipes[0].MASK[0]
                l_hit_mask = self.lower_pipes[0].MASK[0]

                # if bird collided with upipe or lpipe
                u_collide = FlappyBirdEnv.pixel_collision(bird_rect, u_pipe_rect, p_hit_mask, u_hit_mask)
                l_collide = FlappyBirdEnv.pixel_collision(bird_rect, l_pipe_rect, p_hit_mask, l_hit_mask)

                return u_collide or l_collide

    @staticmethod
    def pixel_collision(rect1, rect2, hitmask1, hitmask2):
        """检查两个对象是否发生碰撞，而不仅仅是它们的矩形"""
        rect = rect1.clip(rect2)

        if rect.width == 0 or rect.height == 0:
            return False

        x1, y1 = rect.x - rect1.x, rect.y - rect1.y
        x2, y2 = rect.x - rect2.x, rect.y - rect2.y

        for x in range(rect.width):
            for y in range(rect.height):
                if hitmask1[x1 + x][y1 + y] and hitmask2[x2 + x][y2 + y]:
                    return True
        return False

    @staticmethod
    def rand_action():
        return random.choice((*[FlappyBirdEnv.DO_NOTHING]*6, FlappyBirdEnv.FLAP))

    def close(self):
        pygame.quit()

    def seed(self, seed=None):
        if seed is not None:
            random.seed(seed)
            return [seed]
        else:
            return []

    def render(self, mode='human'):
        pygame.display.update()

    @staticmethod
    def get_action_meanings():
        return ['NOOP', 'FLAP']

    def human_play(self, show_obs=False):
        pause = False
        while True:
            action = FlappyBirdEnv.DO_NOTHING
            for event in pygame.event.get():
                if event.type == pygame.KEYDOWN and event.key == pygame.K_SPACE:
                    action = FlappyBirdEnv.FLAP
                elif event.type == pygame.KEYDOWN and event.key == pygame.K_p:
                    pause = not pause
            if not pause:
                observation, reward, done, info = self.step(action)
                if show_obs:
                    cv2.imshow('obs', observation)
                if done:
                    print('Final score: %d' % self.score.score)
                    pygame.quit()
                    break
            self.render()
